import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from skfem import *
from skfem.models import laplace, mass
from skfem.helpers import dot, grad
import scipy.sparse.linalg as spla
import scipy.sparse as sp

# 参数设置
L = 1.0          # 弦的长度
z = 0.3          # 施加力的位置（距离左端）
Fy = 0.01        # 施加的垂直力大小
T = 10.0         # 弦的张力
rho = 0.1        # 弦的线密度
c = np.sqrt(T/rho)  # 波速
duration = 0.5   # 模拟时长
num_frames = 100 # 动画帧数

# 创建网格
mesh = MeshLine(np.linspace(0, L, 100))
mesh = mesh.with_boundaries({
    'left': lambda x: x[0] == 0,
    'right': lambda x: x[0] == L
})

# 创建有限元空间
element = ElementLineP1()
basis = Basis(mesh, element)

# 组装质量矩阵和刚度矩阵
M = mass.assemble(basis)
K = laplace.assemble(basis) * T

# 手动确定边界点
dofs_list = [0, basis.N - 1]  # 左端点和右端点

# 处理边界条件 - 使用更高效的方法
def apply_boundary_conditions_efficient(A, b=None):
    """应用边界条件到矩阵A和向量b（如果提供），使用更高效的方法"""
    A = A.copy().tolil()  # 转换为LIL格式以高效修改
    
    # 将边界点的行和列设置为单位矩阵
    for dof in dofs_list:
        # 清零行和列
        A[dof, :] = 0
        A[:, dof] = 0
        
        # 设置对角线为1
        A[dof, dof] = 1
    
    A = A.tocsr()  # 转换回CSR格式以提高求解效率
    
    if b is not None:
        b = b.copy()
        # 将边界点的值设置为零
        for dof in dofs_list:
            b[dof] = 0
        return A, b
    return A

# 应用边界条件到刚度矩阵和质量矩阵
K_bc = apply_boundary_conditions_efficient(K)
M_bc = apply_boundary_conditions_efficient(M)

# 初始条件：在z点施加力Fy产生的位移
# 直接创建初始位移，而不是通过求解静态问题
x_points = mesh.p[0]
u0 = np.zeros_like(x_points)

# 找到最接近z的点
idx = np.argmin(np.abs(x_points - z))

# 创建初始位移 - 使用高斯形状而不是点力
sigma = 0.05  # 高斯函数的宽度
u0 = Fy * np.exp(-(x_points - z)**2 / (2 * sigma**2))

# 确保边界点位移为零
for dof in dofs_list:
    u0[dof] = 0

# 初始速度为零
v0 = np.zeros_like(u0)

# 时间步长
dt = 0.0001  # 更小的时间步长
time_steps = np.arange(0, duration, dt)

# 使用中心差分法进行时间积分
u = [u0]
u_prev = u0.copy()  # 前一个时间步的位移

# 预计算系统矩阵
A = M_bc + 0.5 * dt**2 * K_bc

# 时间积分循环
for i in range(1, len(time_steps)):
    # 使用中心差分法
    if i == 1:
        # 第一个时间步使用特殊处理
        b = (2 * M_bc - dt**2 * K_bc) @ u[i-1] - M_bc @ u_prev
    else:
        b = (2 * M_bc - dt**2 * K_bc) @ u[i-1] - A @ u[i-2]
    
    u_new = spla.spsolve(A, b)
    
    # 应用边界条件确保边界点位移为零
    for dof in dofs_list:
        u_new[dof] = 0
    
    u.append(u_new)

# 只取部分时间点用于动画
step = len(u) // num_frames

u_anim = u[::step]
time_anim = time_steps[::step]

# 创建动画
fig, ax = plt.subplots(figsize=(10, 6))
ax.set_xlim(0, L)
ax.set_ylim(-1.5*Fy, 1.5*Fy)
ax.set_xlabel('Position x')
ax.set_ylabel('Displacement u(x,t)')
ax.set_title('String Vibration')
ax.grid(True)

line, = ax.plot([], [], lw=2)
point, = ax.plot([], [], 'ro', ms=8)
time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)

# 初始化函数
def init():
    line.set_data([], [])
    point.set_data([], [])
    time_text.set_text('')
    return line, point, time_text

# 动画更新函数
def animate(i):
    x = mesh.p[0]
    y = u_anim[i]
    line.set_data(x, y)
    point.set_data([z], [0])  # 标记施加力的点
    time_text.set_text(f'Time = {time_anim[i]:.2f} s')
    return line, point, time_text

# 创建动画
ani = FuncAnimation(fig, animate, frames=len(u_anim),
                    init_func=init, blit=True, interval=50)

# 保存为GIF
ani.save('string_vibration.gif', writer='pillow', fps=20)

plt.show()